#include <ntifs.h>
#define NTSTRSAFE_NO_CB_FUNCTIONS
#include <ntstrsafe.h>
#include "../HyperPlatform/kernel_stl.h"
#include "../HyperPlatform/performance.h"
#include "../../Shared/Protocol.h"
#include "NativeEnums.h"
#include "main.h"

extern CProcList *m_IgnoreProcList;
extern CFileList *m_IgnoreFileList;
extern CEventList *m_EventList;

extern PFLT_FILTER m_pFilterHandle;
extern PFLT_PORT m_pClientPort;

NTSTATUS NtFileNameToDosFileName(IN PUNICODE_STRING NtFileName, OUT PUNICODE_STRING DosFileName);

EXTERN_C{

extern DYNAMIC_DATA dynData;

NTSTATUS GetProcessPathByPID(HANDLE ProcessId, PUNICODE_STRING ProcessName);
NTSTATUS GetImagePathByAddress(HANDLE ProcessId, PVOID ImageBase, PUNICODE_STRING ImagePath);

#ifdef _WIN64
NTKERNELAPI PVOID NTAPI PsGetProcessWow64Process(PEPROCESS Process);
#endif

NTKERNELAPI NTSTATUS NTAPI ZwQueryInformationProcess(
	_In_      HANDLE           ProcessHandle,
	_In_      PROCESSINFOCLASS ProcessInformationClass,
	_Out_     PVOID            ProcessInformation,
	_In_      ULONG            ProcessInformationLength,
	_Out_opt_ PULONG           ReturnLength
);

VOID CreateProcessNotifyRoutine(HANDLE ParentId, HANDLE ProcessId, BOOLEAN Create)
{
	if (m_EventList->IsCapturing() && !m_IgnoreProcList->Find(ProcessId))
	{
		UNICODE_STRING ProcessName = { 0 };
		svc_ps_create_process_data *data = (svc_ps_create_process_data *)
			ExAllocatePoolWithTag(PagedPool, sizeof(svc_ps_create_process_data), 'TXSB');

		if (data)
		{
			RtlZeroMemory(data, sizeof(svc_ps_create_process_data));
			data->protocol = svc_ps_create_process;
			data->size = sizeof(svc_ps_create_process_data);
			data->time = PerfGetSystemTime();
			data->Create = Create;
			data->ProcessId = (ULONG)ProcessId;
			data->ParentProcessId = (ULONG)ParentId;

			ProcessName.Buffer = data->ImagePath;
			ProcessName.Length = 0;
			ProcessName.MaximumLength = sizeof(data->ImagePath) - sizeof(WCHAR);

			GetProcessPathByPID(ProcessId, &ProcessName);

			PEPROCESS Process = NULL;
			if (NT_SUCCESS(PsLookupProcessByProcessId(ProcessId, &Process)))
			{
				data->CreateTime = PsGetProcessCreateTimeQuadPart(Process);
				data->SessionId = 0;

				//Get SessionId
				HANDLE ProcessHandle = NULL;
				CLIENT_ID ClientId;
				ClientId.UniqueProcess = ProcessId;
				ClientId.UniqueThread = NULL;
				OBJECT_ATTRIBUTES oa;
				InitializeObjectAttributes(&oa, NULL, OBJ_KERNEL_HANDLE, 0, 0);
				if (NT_SUCCESS(ZwOpenProcess(&ProcessHandle, PROCESS_QUERY_INFORMATION, &oa, &ClientId)))
				{
					PROCESS_SESSION_INFORMATION psi;
					if (NT_SUCCESS(ZwQueryInformationProcess(ProcessHandle, ProcessSessionInformation, &psi, sizeof(psi), NULL)))
					{
						data->SessionId = psi.SessionId;
					}
					ZwClose(ProcessHandle);
				}

#ifdef _WIN64
				data->Is64Bit = (PsGetProcessWow64Process(Process) == NULL) ? TRUE : FALSE;
#else
				data->Is64Bit = FALSE;
#endif
				ObDereferenceObject(Process);
			}

			m_EventList->SendEvent(data);
		}
	}

	if(!Create)
		m_IgnoreProcList->RemoveProcess(ProcessId);
}

typedef struct
{
	PFILE_OBJECT FileObject;
	PFLT_FILE_NAME_INFORMATION pFileInfo;
	PKEVENT NotifyEvent;
	NTSTATUS Status;
}QueryFileName_t;

VOID QueryFileDosNameRoutine(QueryFileName_t *Parameter)
{
	PFLT_FILE_NAME_INFORMATION pFileInfo;
	Parameter->Status = FltGetFileNameInformationUnsafe(Parameter->FileObject, NULL, FLT_FILE_NAME_NORMALIZED | FLT_FILE_NAME_QUERY_DEFAULT, &pFileInfo); //IoQueryFileDosDeviceName(Parameter->FileObject, &pNameInfo);
	if (NT_SUCCESS(Parameter->Status))
	{
		Parameter->pFileInfo = pFileInfo;
	}
	KeSetEvent(Parameter->NotifyEvent, 0, FALSE);
}

VOID LoadImageNotifyCallback(PUNICODE_STRING FullImageName, HANDLE ProcessId, PIMAGE_INFO pImageInfo)
{
	LARGE_INTEGER time = { 0 };
	
	if (m_IgnoreProcList->Find(ProcessId))
		return;

	if (!m_EventList->IsCapturing())
		return;

	svc_ps_load_image_data *data = (svc_ps_load_image_data *)
		ExAllocatePoolWithTag(PagedPool, sizeof(svc_ps_load_image_data), 'TXSB');

	if (data)
	{
		RtlZeroMemory(data, sizeof(svc_ps_load_image_data));

		data->protocol = svc_ps_load_image;
		data->size = sizeof(svc_ps_load_image_data);
		data->time = PerfGetSystemTime();
		data->ProcessId = (ULONG)ProcessId;
		if(pImageInfo->SystemModeImage)
			data->ProcessId = 4;
		data->ImageBase = (ULONG64)pImageInfo->ImageBase;
		data->ImageSize = (ULONG64)pImageInfo->ImageSize;

		UNICODE_STRING ImageFileName;
		ImageFileName.Buffer = data->ImagePath;
		ImageFileName.Length = 0;
		ImageFileName.MaximumLength = sizeof(data->ImagePath) - sizeof(WCHAR);

		if (pImageInfo->ExtendedInfoPresent)
		{
			KEVENT QueryEvent;
			WORK_QUEUE_ITEM QueryWorkItem;
			QueryFileName_t QueryParam = { 0 };
			PIMAGE_INFO_EX pImageInfoEx = CONTAINING_RECORD(pImageInfo, IMAGE_INFO_EX, ImageInfo);

			QueryParam.FileObject = pImageInfoEx->FileObject;
			QueryParam.NotifyEvent = &QueryEvent;

			KeInitializeEvent(&QueryEvent, NotificationEvent, FALSE);
			ExInitializeWorkItem(&QueryWorkItem, (PWORKER_THREAD_ROUTINE)QueryFileDosNameRoutine, &QueryParam);
			ExQueueWorkItem(&QueryWorkItem, DelayedWorkQueue);			
			KeWaitForSingleObject(&QueryEvent, Executive, KernelMode, FALSE, 0);

			if (NT_SUCCESS(QueryParam.Status) && QueryParam.pFileInfo != NULL)
			{
				if (!NT_SUCCESS(NtFileNameToDosFileName(&QueryParam.pFileInfo->Name, &ImageFileName)))
					RtlCopyUnicodeString(&ImageFileName, &QueryParam.pFileInfo->Name);

				FltReleaseFileNameInformation(QueryParam.pFileInfo);
			}
			else
			{
				if(ARGUMENT_PRESENT(FullImageName))
					RtlCopyUnicodeString(&ImageFileName, FullImageName);
			}
		}
		else
		{
			if (ARGUMENT_PRESENT(FullImageName))
				RtlCopyUnicodeString(&ImageFileName, FullImageName);
		}

		m_EventList->SendEvent(data);
	}
}

}